from langchain.llms.base import LLM
from typing import Optional, List, Any, Mapping
from langchain.callbacks.manager import CallbackManagerForLLMRun
from openai import OpenAI
from os import getenv
import os
from config.config import Config

os.environ['OPENROUTER_API_KEY'] = Config.OPENROUTER_API_KEY
class ProxyGPT(LLM):
    model:str = 'openai/gpt-4o'

    @property
    def _llm_type(self) -> str:
        return self.model

    @property
    def _llm_type(self) -> str:
        return self.model

    def _call(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> str:
        # gets API Key from environment variable OPENAI_API_KEY
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=getenv("OPENROUTER_API_KEY"),
        )

        completion = client.chat.completions.create(
            extra_headers={
            },
            model="openai/gpt-4o",
            messages=[
                {
                    "role": "user",
                    "content": prompt,
                },
            ],
        )
        response = completion.choices[0].message.content
        return response

    def chat_history_by_list(messagesList: list) -> str:
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=getenv("OPENROUTER_API_KEY"),
        )

        completion = client.chat.completions.create(
            extra_headers={
            },
            model="openai/gpt-4o",
            messages = messagesList,
        )
        response = completion.choices[0].message.content
        return response
